import os
os.chdir("/home/user/zsy/rag-demo/compute")
from compute_utils import *
import json
import copy
import re

class Shared:
    with open('expressions.json', 'r', encoding='utf-8') as f:
        expresstions = json.load(f)["compute"]
    with open('get_people.json', 'r', encoding='utf-8') as f:
        get_people = json.load(f)
    oprators = ['+', '-', '*', '/']

class Compute:
    def __init__(self, ex_key = ""):
        if ex_key.__contains__("城镇居民生活用水定额"):
            pass
        self.expresstions = Shared.expresstions
        self.values = Shared.get_people["get"]
        self.values.update(Shared.get_people["people"])
        self.original_key = ex_key
        self.ex_key = self.original_key #存放公式名
        self.history = [] #存储复合公式的子公式[[公式名，公式],...]
        self.key_value = dict() #存储子公式计算结果
        
        def process_key(keys): #提取关键词（年，市区等）并让self.ex_key为剩余的内容
            key = ""
            for x in keys:
                if self.ex_key.find(x) != -1:
                    key = x
                    break
            if key == "":
                result = ""
            else:
                result = self.ex_key[:self.ex_key.find(key)+len(key)]
                self.ex_key = self.ex_key[self.ex_key.find(key)+len(key):]
            return result
        
        self.year = process_key(["年"])
        self.city = process_key(["市区","中心","市","县"]) #处理 市 区 县
        
    def process_expresstion(self, ex_list): #计算核心，负责替换公式中的变量，最后返回计算结果
        unit = ex_list[-1]
        ex_list = ex_list[:-1]
        result = "" #存储纯数值表达式
            
        for x in ex_list:
            if len(x) == 0: # 忽略空字符串，因为有些公式没有完善
                continue
            if x.__contains__("城镇居民生活用水定额"):
                pass
            x = x.replace("@", self.year)
            x = x.replace("#", self.city)
            if x[0] == '!': #检查改公式是否需要嵌套一层
                t_cp = Compute(x[1:])
                t_cp.compute()
                try:
                    if isinstance(t_cp.exp_result, str):#当返回结果为str则认为出错
                        return t_cp.exp_result
                    single = [t_cp.exp_result[2] + t_cp.exp_result[-1], t_cp.ex_key]
                except Exception as e:
                    print(e)
                    print("Key Error:",t_cp.original_key)
                    # exit()
            else:
                single = search_key(self.values, self.year, self.city, x) #single [带符号的数值, 公式名]
            if single is not None: #跳过特殊字符（加减乘除等），并替换为values中的值
                match = re.match(r"([0-9.]+)(.*)", single[0]) #group(1)是数值，group(2)是单位
                result += match.group(1)
                self.key_value[single[1]] = match.group(1)
            else:
                result += x
        try:
            return [''.join(ex_list), result, eval(''.join(result)), unit]
        except:
            # print(f"计算错误：\n{''.join(ex_list)}\n->{result}","\n符号位:",unit) #决定报错时是否结束程序
            # exit()
            return f"计算错误：\n{''.join(ex_list)}\n->{result}"
    
    def compute(self): #对复合公式处理，然后直接调用计算核心
        if self.original_key.__contains__("取水量"):
            pass
        res = search_key(self.expresstions, self.year, self.city, self.ex_key) #不要第一个值，只用来测试各种组合后，是否存在expresstions中
        if self.original_key in self.values:
            self.exp_result = ["","",str(self.values[self.original_key]),""] #用于递归时访问，模拟返回的列表值
            return str(self.values[self.original_key]) #当get_people.json中存在，则直接返回
        elif res is not None:
            self.ex_key = res[1]
            self.expresstion = str_to_expression(self.expresstions[self.ex_key]) #在公式表中找到的目标公式，最后一个是符号位
        else:
            self.exp_result = f"缺少公式:{self.ex_key}"
            return "找不到该公式"
        
        i = 0
        while i < len(self.expresstion):#首先将复合公式拆解
            if self.expresstion[i] in self.expresstions:
                self.history.insert(0,[self.expresstion[i], str_to_expression(self.expresstions[self.expresstion[i]])]) #每拆解一次，就存储为一个公式，用于模拟计算过程
                self.expresstion = insert_list(self.expresstion, str_to_expression(self.expresstions[self.expresstion[i]])[:-1], i)
                i = 0
            else:
                i += 1
        if len(self.history) > 0:#先计算子公式结果
            for x in self.history:
                result = self.process_expresstion(x[1])
                if isinstance(result[2], float):
                    self.values[x[0]] = result[2]
                    self.key_value[x[0]] = f'{result[2]:.2f}' + result[-1]
    
        self.exp_result = self.process_expresstion(self.expresstion) #计算主公式的结果
        if isinstance(self.exp_result,str):
            return self.exp_result
        else:
            self.exp_result[2] = f'{self.exp_result[2]:.4f}' #要精确的位数，get_pople.json中默认两位
            self.key_value[self.ex_key] = self.exp_result[2] + self.exp_result[-1] #将最后计算的结果保存到key_value中
            return "计算过程:\n" + '\n->'.join(self.exp_result[:-1]) + self.exp_result[-1]
        
def process_file(file_name = "input.txt", output_file_name = "output.json"): #读取txt文件作为输入list，输出output.json格式
    all_data_dict = dict()
    exp_list = read_txt(file_name)
    for x in exp_list:
        x = x.strip()
        x_ori = x
        if x == "规划水平年水厂自用水量":
            pass
        x = x.replace("城区","")
        x = x.replace("城城","城") #**县城 城区... 改为 **县 城区...
        cp = Compute(x)
        res = cp.compute()
        if not res.startswith("计算过程"):
            all_data_dict[x_ori] = res
        else:
            if len(cp.history)>0:
                res += "\n子公式计算:\n"
                for x in cp.history:
                    res += f"{x[0]}: {x[1]}\n->{cp.key_value[x[0]]}\n"
            res += "参考值:\n"
            for x in cp.key_value.keys():
                if x not in cp.history:
                    res += f"{x}: {cp.key_value[x]}\n"
            all_data_dict[x_ori] = res
    with open(output_file_name, "w") as f:
      json.dump(all_data_dict, f, ensure_ascii=False, indent=4)
      
def deep_dict(body, h_key): #递归处理字典
    for k, v in body.items():
        if isinstance(v, dict):
            deep_dict(v, h_key + k)
        else:
            if any(x in Shared.oprators for x in v): #考虑value是公式的情况
                cp = Compute(k)
                cp.ex_key = k
                cp.values.update(Shared.count_history) #将表中计算的值更新到常量表中
                cp.expresstions = body
                result = cp.compute()
            else:
                cp = Compute(v)
                result = cp.compute()
            Shared.count_history[h_key + k] = cp.exp_result[2]
            if Shared.detailed: #判断是否输出计算细节
                if result.startswith("计算过程"):
                    body[k] = result
                else:
                    body[k] = "计算过程:\n" + body[k] + "\n->" + result
            else:
                body[k] = cp.exp_result[2] + cp.exp_result[-1]
    
def process_json(file_name = "input.json", output_file_name = "output_table"):
    exp_table = read_json(file_name)
    Shared.result_dict = copy.deepcopy(exp_table) #输出的细节版本用该字典存储
    Shared.count_history = dict() #计算结果存储，用于表格内公式计算
    for k in exp_table.keys():
        Shared.detailed = False #计算不输出计算细节的版本
        deep_dict(Shared.result_dict[k], k)
        Shared.detailed = True #计算输出计算细节的版本
        deep_dict(exp_table[k], k)
    with open(output_file_name + ".json", 'w', encoding='utf-8') as f:
        json.dump(Shared.result_dict, f, ensure_ascii=False, indent=4)
    with open(output_file_name + "_detailed.json", 'w', encoding='utf-8') as f:
        json.dump(exp_table, f, ensure_ascii=False, indent=4)
        
    
if __name__ == "__main__":
    #使用说明：查询语句需要满足：**["年"]**["市区","中心","市","县"]+公式名称；公式前为1部分，公式为2部分。
    #年可以是规划水平年，水平年等以年结尾，2部分需要在expressions.json中存在。
    #2部分中的value参数需要存在get_people.json中存在。
    #2部分的最后一个参数为单位。
    #表格引用方式：一级一级拼接字符即可实现引用，涉及到计算符号的value才需要单位，否则不可在末尾加入单位
    #             value中如果既用了计算又想要引用公式计算，需要在公式前加入!符号
    #符号介绍：!表示嵌套计算，#表示市区替换，@表示年替换
    #公式编写踩坑：1.注意符号位是否存在 2.注意公式编写是否正确 3.公式引用的值是否存在
    #             4.是否需要递归一层计算（跨年或跨市还不是可引用的常数）
    # ex_key = "规划水平年用水户耗水量"
    # cp = Compute(ex_key)
    # result = cp.compute()
    # print(result)
    
    # process_file()
    process_json()
